# coding = utf-8

import numpy as np
from metrics import r2_score


class SimpleLinearRegression:

    def __init__(self):
        self.a_ = None
        self.b_ = None

    def fit(self, x, y):
        assert x.shape[0] == y.shape[0], "size must be same"

        x_mean = np.mean(x)
        y_mean = np.mean(y)

        num = (x - x_mean).dot(y - y_mean)
        d = (x - x_mean).dot(x - x_mean)
        self.a_ = num / d
        self.b_ = y_mean - self.a_ * x_mean

        return self

    def predict(self, x_predict):
        assert self.a_ is not None and self.b_ is not None, "must fit before predict"

        return np.array([self._predict(x_i) for x_i in x_predict])

    def _predict(self, x_single):
        return self.a_ * x_single + self.b_

    def score(self, y_true, y_predict):
        return r2_score(y_true, y_predict)

    def __repr__(self):
        return "SimpleLinearRegression()"
